package com.huatai.datacenter.service;

import com.huatai.datacenter.config.TopicConfig;
import com.huatai.datacenter.entity.KafkaRecord;
import com.huatai.datacenter.utils.KafkaConsumers;
import com.huatai.datacenter.websocket.WebSocketServer;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.time.Duration;
import java.util.*;

import static org.apache.kafka.clients.producer.ProducerConfig.*;

/**
 * @author Lion
 * @date 2023/3/27  9:24
 */
@Slf4j
@Service
public class KafkaConsumerService {

	@Autowired
	private ClusterService clusterService;

	/**
	 * @author Lion
	 * @description 普通方式消费Topic
	 * @param clusterId 集群id
	 * @param groupName 分组名称
	 * @param topicName 组题名称
	 * @param dataSize 一次消费多少条数据
	 * @param duration 时间
	 * @param isCommit 是否提交偏移量
	 * @return org.apache.kafka.clients.consumer.ConsumerRecords<java.lang.String,java.lang.String>
	 */
	public ConsumerRecords<String,String> consumer(String clusterId, String groupName, String topicName, int dataSize, Duration duration,boolean isCommit){

		ConsumerRecords<String,String> result = null;
		//获取集群
		String broker = clusterService.getBrokers(clusterId);
		//配置消费者信息
		Properties consumerProps =generatorConsumerProps(broker,groupName,isCommit);
		consumerProps.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG,dataSize);
		KafkaConsumers<String,String> consumers = new KafkaConsumers<>(consumerProps);
		consumers.subscribe(topicName);
		try {
			result = consumers.poll(duration);
		}catch (Exception e){
			log.error("  consumer  topic: {} error! ", topicName, e);
			throw e;
		}finally {
			try {
				if (isCommit){
					consumers.commit();
				}
				consumers.close();
			} catch (IOException e) {
				log.error("close consumer error! ");
			}
		}
		return result;
	}

	/**
	 * @author Lion
	 * @description 通过指定partition和offset方式消费topic
	 * @param clusterId 集群id
	 * @param groupName 分组名称
	 * @param topicName 组题名称
	 * @param dataSize 一次消费多少条数据
	 * @param duration 时间周期
	 * @param isCommit 是否提交偏移量
	 * @param partition 分区
	 * @param offset 偏移量
	 * @return org.apache.kafka.clients.consumer.ConsumerRecords<java.lang.String,java.lang.String>
	 */
	public ConsumerRecords<String,String> consumer(String clusterId, String groupName, String topicName, int dataSize, Duration duration,boolean isCommit,Integer partition,Long offset){
		ConsumerRecords<String,String> result = null;
		Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
		//获取集群
		String broker = clusterService.getBrokers(clusterId);
		Properties consumerProps = generatorConsumerProps(broker,groupName,false);
		consumerProps.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG,dataSize);
		KafkaConsumers<String,String> consumer = new KafkaConsumers<String, String>(consumerProps);
		TopicPartition topicPartition = new TopicPartition(topicName,partition);
		consumer.assign(topicPartition);
		consumer.seek(topicPartition,offset);

		try {
			result = consumer.poll(duration);
			List<ConsumerRecord<String,String>> records = result.records(topicPartition);
			if (!records.isEmpty()){
				long lastOffset = records.get(records.size() - 1).offset();
				offsets.put(topicPartition,new OffsetAndMetadata(lastOffset + 1));
			}

		}catch (Exception e){
			log.error("consumer topic: {} error! ", topicName, e);
			throw e;
		}finally {
			if (isCommit){
				consumer.commitByPartition(offsets);
			}
			try {
				consumer.close();
			} catch (IOException e) {
				log.error("close consumer error!");
			}
		}
		return result;
	}

	public void consumerDD(String clusterId, String groupName, String topicName, int dataSize, Duration duration,boolean isCommit){

		ConsumerRecords<String,String> result = null;
		//获取集群
		String broker = clusterService.getBrokers(clusterId);
		//配置消费者信息
		Properties consumerProps =generatorConsumerProps(broker,groupName,isCommit);
		consumerProps.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG,dataSize);
		KafkaConsumers<String,String> consumers = new KafkaConsumers<>(consumerProps);
		consumers.subscribe(topicName);

		while (true){
			try {
				result = consumers.poll(duration);

				for (ConsumerRecord<String, String> record : result){
					KafkaRecord kafkaRecord = new KafkaRecord();
					kafkaRecord.setPartition(record.partition());
					kafkaRecord.setOffset(record.offset());
					kafkaRecord.setKey(record.key());
					log.info(record.key());
					kafkaRecord.setValue(record.value());
					log.info(record.value());
				}
			}catch (Exception e){
				log.error("  consumer  topic: {} error! ", topicName, e);
				throw e;
			}
		}
	}

	/**
	 * @author Lion
	 * @description 根据分组和主题名称获取数据大小
	 * @param brokers
	 * @param groupName
	 * @param topicName
	 * @return java.util.Map<org.apache.kafka.common.TopicPartition,java.lang.Long>
	 */
	public Map<TopicPartition,Long> getLogSize(String brokers,String groupName,String topicName){

		KafkaConsumer<String,String> consumer = null;
		Map<TopicPartition,Long> result = new HashMap<>();
		List<TopicPartition> topicPartitions = new ArrayList<>();

		try {
			consumer = generatorKafkaCustomer(brokers,groupName,false);

			List<PartitionInfo> partitionInfos = consumer.partitionsFor(topicName);
			if (partitionInfos == null){
				return result;
			}
			partitionInfos.forEach(partition ->{
				TopicPartition topicPartition = new TopicPartition(topicName,partition.partition());
				topicPartitions.add(topicPartition);
			});
			result = consumer.endOffsets(topicPartitions);
		}catch (Exception e){
			log.error("consumer getLogSize error.", e);
		}finally {
			consumer.close();
		}
		return result;
	}

	/**
	 * @author Lion
	 * @description 提交偏移量
	 * @param clusterId
	 * @param group
	 * @param topic
	 * @return boolean
	 */
	public boolean commitOffsetLatest(String clusterId,String group,String topic){

		String brokers = clusterService.getBrokers(clusterId);
		KafkaConsumer<String,String> consumer = generatorKafkaCustomer(brokers,group,true);
		try {
			List<PartitionInfo> partitionInfos = consumer.partitionsFor(topic);
			List<TopicPartition> topicPartitions = new ArrayList<>();
			partitionInfos.forEach(topicPartition -> {
				TopicPartition partition = new TopicPartition(topic,topicPartition.partition());
				topicPartitions.add(partition);
			});
			consumer.assign(topicPartitions);
			Map<TopicPartition,Long> result = consumer.endOffsets(topicPartitions);
			result.forEach(consumer :: seek);
			return true;
		}catch (Exception e){
			log.error("consumer commitOffsetLatest error.", e);
			return false;
		}finally {
			consumer.close();
		}
	}

	/**
	 * @author Lion
	 * @description 生成kafka客户端信息
	 * @param brokers
	 * @param group
	 * @param isAutoCommit
	 * @return org.apache.kafka.clients.consumer.KafkaConsumer<java.lang.String,java.lang.String>
	 */
	private KafkaConsumer<String,String> generatorKafkaCustomer(String brokers,String group,boolean isAutoCommit){
		Properties properties = generatorConsumerProps(brokers,group,isAutoCommit);
		KafkaConsumer<String,String> consumer = new KafkaConsumer<String, String>(properties);
		return consumer;
	}

	/**
	 * @author Lion
	 * @description kafka配置信息
	 * @param broker
	 * @param group
	 * @param isAutoCommit
	 * @return java.util.Properties
	 */
	public Properties generatorConsumerProps(String broker, String group, boolean isAutoCommit){
		Properties consumerProps = new Properties();
		consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,broker);
		consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG,group);
		consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG,isAutoCommit);
		consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
		consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,StringDeserializer.class);
		consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, TopicConfig.EARLIEST);

		return consumerProps;
	}

	public Properties generatorProducerProps(String broker){
		Properties producerProps = new Properties();
		producerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,broker);
		producerProps.put(ACKS_CONFIG, "all");
		producerProps.put(KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringSerializer");
		producerProps.put(VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringSerializer");

		return producerProps;
	}


}
